#
# messkit.session
#
# Copyright 2003-2009 Helsinki Institute for Information Technology (HIIT)
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""
Bidirectional message pipe over TCP or other byte stream
"""

import time
import socket
import logging

from messkit.nson import encode, decode
from messkit.netstring import parse_netstring

class SessionStats:
    """
    Blackboard for session statistics

    Each session has an instance of this in the attribute "stats".
    """
    def __init__(self):
        self.packets_in = 0
        self.packets_out = 0
        self.bytes_in = 0
        self.bytes_out = 0
        self.start_timestamp = time.time()
        self.in_timestamp = self.start_timestamp
        self.out_timestamp = self.start_timestamp

class Session(object):
    """
    Bidirectional message pipe over TCP or other byte stream

    All methods are nonblocking.

    All callbacks (startup, shutdown, and message handlers) are
    invoked asynchronously in the main thread, using the stream
    manager's callback callback mechanism if necessary.
    """
    def __init__(self, stream, startup=None, shutdown=None,
                 handler_prefix=None, **keys):
        """
        Initialize a communication session.

        The first argument must be a Stream instance, which will
        subsequently be managed by the session.  If it is an
        initiating endpoint, it will be connected when the session's
        connect() method is called.

        Startup and shutdown callbacks can be specified with keyword
        arguments.  Any unrecognized keyword arguments are passed
        along to the startup callback.

            startup(error, **keys)
            shutdown()

        If callbacks are not specified, the startup() and shutdown()
        methods of the session are called instead.

        The startup callback is called at both endpoints after the
        stream has been connected and before any other callbacks.  The
        shutdown callback is called at both endpoints after the stream
        has been closed, and it is always the last callback.

        The value passed to startup as "error" will normally be None.
        However, if this is the initiating endpoint and the connection
        attempt fails, the startup callback will be invoked with a
        value that tests true and indicates the reason.  The shutdown
        callback will be invoked immediately thereafter.

        A session may also define message handlers using register_function().
        A handy place to do this is in the startup handler.

        If specified, the handler prefix determines the mapping
        assumed by register_function() between message method names
        and handler function names.  If the prefix is, for instance,
        "handle_", register_function() will assume that a function
        named "handle_foo" is intended to handle messages with method
        name "foo".

        A connection can be closed for any number of reasons:

         * The application has called the close() method.

         * A handler function (the startup callback or a message
           handler) has raised an exception.

         * A protocol error has been detected.  These include badly
           formatted messages and messages for which no handler has
           been defined.

         * The application at the other end has closed the connection
           for one of the above reasons.  These cases are
           indistinguishable at this end.

         * A send or receive socket operation has raised an exception.
           In this case, self.stream.error will contain the exception
           object.

         * The connection attempt has failed (assuming this is the
           initiating endpoint).  In this case, self.stream.error will
           contain the exception object, and this will have been
           passed to the startup callback.
        """
        self.stream = stream
        self.startup_callback = startup or self.startup
        self.shutdown_callback = shutdown or self.shutdown
        self.keyword_args = keys

        self.funcs = {}
        self.handler_prefix = handler_prefix
        self.stats = SessionStats()
        self.input_suspended = False

        self.__description = None

        if stream.connected:
            self.issue_callback(self.__startup, None)

    @property
    def closed(self):
        return self.stream.closed

    def register_function(self, function, name=None):
        """
        Define a handler for incoming messages.

        This is modeled after Python's SimpleXMLRPCServer.

        If "name" is given, it will be the method name associated with
        the function.  Otherwise, function.__name__ will be used,
        after stripping off the handler prefix if one has been set.

        The name can be either a normal or Unicode string, and may
        contain characters not legal in Python identifiers, including
        the period character.

        The function is called as function(*params, **options), where
        the params and options are as provided to send() at the other
        endpoint.
        """
        if name is None:
            name = self._get_method_name(function)

        self.funcs[name] = function

    def unregister_function(self, function, name=None):
        """
        Remove a handler for incoming messages.

        The arguments are as for register_function().  The function
        and method name must match a registered handler.
        """
        if name is None:
            name = self._get_method_name(function)

        assert self.funcs[name] == function
        del self.funcs[name]

    def connect(self, address):
        """
        Connect the underlying stream to a remote address.

        This is only used at the initating endpoint.  The argument is
        a tuple (host, port) specifying the remote address to connect to.

        This always returns immediately.  The startup callback will be
        invoked upon completion.
        """
        self.stream.connect(address, self.__startup)

    def handshake(self, cb=None, role=None, **params):
        """
        Optionally perform a certificate-based TLS handshake.

        This can only be called from the startup handler or from a
        message handler.  If the handshake completes successfully, the
        callback will be invoked with the hash of the other endpoint's
        public key as its argument.

        The "role" can be either "client" or "server".  By default,
        the endpoint that initiated the TCP connection will play the
        role of client.  If messages are exchanged before the TLS
        handshake is initiated, the handshake must be done in the
        opposite direction of the last message exchanged, and such
        that one can be sure that no messages are still in flight.

        Additional keyword parameters relate to the details of TLS:

        certChain - the local entity's CryptoID certificate chain (an
        instance of cryptoIDlib.CertChain.CertChain, initialized by
        passing an XML certificate chain to the parse() method).

        privateKey - the local entity's RSA key pair, as returned by
        tlslite.utils.keyfactory.parseXMLKey(str, private=True).

        settings (optional) - an instance of tlslite.api.HandshakeSettings.
        """
        def proxy(error):       # Stream callback
            if not error:
                self._call(cb, self.stream.end_entity)
                self.input_suspended = False
                self.__process_input() # Restart processing.
            else:
                self.close(flush=False)

        self.input_suspended = True
        self.stream.handshake(proxy, role, **params)

    def send(self, method, *params, **options):
        """
        Send a one-way message to the other endpoint.

        This queues the message for transmission and returns
        immediately.  It has no effect if the session has been closed.
        """
        if not self.closed:
            payload = {'method': method}
            if params:
                payload['params'] = params
            if options:
                payload['options'] = options

            s = encode(payload)

            self.stats.packets_out += 1
            self.stats.bytes_out += len(s)
            self.stats.out_timestamp = time.time()

            self.stream.write(s)
            self.stream.flush()

    def close(self, flush=True):
        """
        Close the communication session.

        Calling close() will close the session at both ends.  It can
        be called more than once.

        This closes the session abruptly.  A best-effort attempt is
        made to send buffered outbound messages, but buffered inbound
        messages will not be received.
        """
        self.stream.close(flush, self.__shutdown)

    def startup(self, error):
        """
        Default startup handler
        """
        pass

    def shutdown(self):
        """
        Default shutdown handler
        """
        pass

    def issue_callback(self, f, *args):
        self.stream.issue_callback(f, *args)

    def __shutdown(self):       # Stream callback
        self._call(self.shutdown_callback)

        del self.funcs
        del self.startup_callback
        del self.shutdown_callback

    def __startup(self, error): # Stream callback
        self._call(self.startup_callback, error, **self.keyword_args)

        if error:
            self.close(flush=False) # Connect failed.
        else:
            self.stream.set_error_handler(self._handle_stream_error)
            self.__process_input()  # Start processing.

    def __process_input(self):
        try:
            stream = self.stream
            while not self.closed and not self.input_suspended:
                s = stream.peek(20)
                content, k = parse_netstring(s)
                if not k:
                    stream.fill(self.__receive)
                    return
                elif k > len(stream):
                    stream.fill(self.__receive, size=k)
                    return
                else:
                    s = stream.read(k)

                    self.stats.packets_in += 1
                    self.stats.bytes_in += len(s)
                    self.stats.in_timestamp = time.time()

                    payload = decode(s)
                    method = payload['method']
                    params = payload.get('params', [])
                    options = payload.get('options', {})
                    self._dispatch(method, params, options)
        except Exception:
            self._handle_protocol_error()

    def __receive(self, n):     # Stream callback
        try:
            if n > 0:
                self.__process_input() # Continue processing.
            elif len(self.stream) == 0:
                self.close(flush=True) # Normal EOF.
            else:
                raise ValueError("Received truncated message.")
        except Exception:
            self._handle_protocol_error()

    def _handle_stream_error(self, e):
        logging.warning('Stream error: %s', e)
        self.close(flush=False)

    def _handle_protocol_error(self):
        logging.exception("Protocol error:")
        self.close(flush=False)

    def _handle_handler_error(self):
        # This probably indicates a bug in the application program.
        logging.exception("Error in handler function:")
        self.close(flush=False)

    def _dispatch(self, method, params, options):
        """
        Dispatch an incoming message to a handler function.

        This implements the mapping to handlers registered by
        register_function().  It can be useful to override this
        method.
        """
        f = self.funcs.get(method)
        if f is not None:
            self._call(f, *params, **options)
        else:
            raise ValueError('Method "%s" is not supported.' % method)

    def _call(self, f, *params, **options):
        if f is not None:
            try:
                f(*params, **options)
            except Exception:
                self._handle_handler_error()

    def _get_method_name(self, function):
        name = function.__name__
        prefix = self.handler_prefix
        if prefix is not None:
            if name.startswith(prefix):
                name = name[len(prefix):]
            else:
                raise ValueError(
                    'Function name "%s" does not start with "%s".'
                    % (name, prefix))
        return name

    def get_description(self, lookup_hostname=True):
        address = self.stream.address
        if address is None:
            return 'Uninitialized connection'

        if self.__description is None:
            host, port = address
            if lookup_hostname:
                host = socket.getfqdn(host)

            if self.stream.initiating:
                self.__description = 'Connection to %s port %s' % (host, port)
            else:
                self.__description = 'Connection from %s' % host

        return self.__description
